"""
StarGAN v2
Copyright (c) 2020-present NAVER Corp.

This work is licensed under the Creative Commons Attribution-NonCommercial
4.0 International License. To view a copy of this license, visit
http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
"""

import os
import paddle
import pickle
import six


class CheckpointIO(object):
    def __init__(self, fname_template, **kwargs):
        os.makedirs(os.path.dirname(fname_template), exist_ok=True)
        self.fname_template = fname_template
        self.module_dict = kwargs

    def register(self, **kwargs):
        self.module_dict.update(kwargs)

    def save(self, step):
        fname = self.fname_template.format(step)
        print('Saving checkpoint into %s...' % fname)
        outdict = {}
        for name, module in self.module_dict.items():
            outdict[name] = module.state_dict()
        self._save(outdict, fname)

    def load(self, step):
        fname = self.fname_template.format(step)
        assert os.path.exists(fname), fname + ' does not exist!'
        print('Loading checkpoint from %s...' % fname)
        module_dict = self._load(fname)
        for name, module in self.module_dict.items():
            module.set_state_dict(module_dict[name])

    def _save(self, state_dicts, file_name):
        def convert(state_dict):
            model_dict = {}

            for k, v in state_dict.items():
                if isinstance(
                        v,
                    (paddle.fluid.framework.Variable, paddle.fluid.core.VarBase)):
                    model_dict[k] = v.numpy()
                else:
                    model_dict[k] = v

            return model_dict

        final_dict = {}
        for k, v in state_dicts.items():
            if isinstance(
                    v,
                (paddle.fluid.framework.Variable, paddle.fluid.core.VarBase)):
                final_dict = convert(state_dicts)
                break
            elif isinstance(v, dict):
                final_dict[k] = convert(v)
            else:
                final_dict[k] = v

        with open(file_name, 'wb') as f:
            pickle.dump(final_dict, f, protocol=2)

    def _load(self, file_name):
        with open(file_name, 'rb') as f:
            state_dicts = pickle.load(f) if six.PY2 else pickle.load(
                f, encoding='latin1')
        return state_dicts